import os
import uuid
import yaml
import sqlite3
import psycopg2
from psycopg2.extras import DictCursor
from pgvector.psycopg2 import register_vector
from contextlib import contextmanager
from loguru import logger
from dotenv import load_dotenv
import tempfile
import json
import numpy as np
from typing import List, Dict, Any, Optional, Union
from psycopg2.extensions import connection as Psycopg2Connection
from sqlite3 import Connection as SQLiteConnection

class DatabaseManager:
    """数据库管理器，处理所有数据库操作和向量功能"""
    
    _instance = None
    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(DatabaseManager, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if self._initialized:
            return
        
        # 加载.env文件
        load_dotenv()
        
        self.connection = None
        self.in_memory_mode = False
        self.vector_dimension = 1536  # 默认向量维度
        
        # 实现PostgreSQL优先，SQLite备选的策略
        postgres_dsn = None
        sqlite_dsn = 'sqlite:///./data/aichat.db'  # 确保使用./data/aichat.db
        
        # 尝试从环境变量和配置中获取PostgreSQL连接信息
        try:
            # 优先从DATABASE_URL获取PostgreSQL连接
            postgres_dsn = os.getenv('DATABASE_URL')
            if postgres_dsn and not postgres_dsn.startswith('postgres'):
                # 如果提供了URL但不是PostgreSQL，忽略它
                postgres_dsn = None
            
            # 如果DATABASE_URL不存在或不是PostgreSQL，尝试从配置中获取
            if not postgres_dsn:
                try:
                    from backend.common.config import DatabaseConfig
                    db_config_obj = DatabaseConfig()
                    if db_config_obj.dsn and db_config_obj.dsn.startswith('postgres'):
                        postgres_dsn = db_config_obj.dsn
                        self.vector_dimension = getattr(db_config_obj, 'vector_dimension', 1536)
                except Exception:
                    pass
        except Exception as e:
            logger.warning(f"获取数据库配置时出错: {str(e)}")
        
        # 先尝试使用PostgreSQL
        if postgres_dsn:
            logger.info(f"尝试使用PostgreSQL数据库: {self._mask_dsn(postgres_dsn)}")
            self.dsn = postgres_dsn
            self.db_type = 'postgres'
        else:
            # 如果没有PostgreSQL配置，使用SQLite
            logger.info(f"未找到PostgreSQL配置，使用SQLite数据库")
            self.dsn = sqlite_dsn
            self.db_type = 'sqlite'
        
        # 初始化数据库连接
        self._connect()
        # 初始化表结构
        self._initialize_tables()
        
        # 内存模式数据存储
        if self.in_memory_mode:
            self.memory_vectors = []
            self.next_vector_id = 1
        
        self._initialized = True
    
    def _mask_dsn(self, dsn):
        """隐藏DSN中的敏感信息"""
        if '@' in dsn:
            parts = dsn.split('@')
            return f"{parts[0].split(':')[0]}:***@{parts[1]}"
        return dsn
        
    def _connect(self):
        """建立数据库连接"""
        try:
            if self.db_type == 'sqlite':
                # 从DSN中提取路径，格式: sqlite:///path/to/db
                if self.dsn.startswith('sqlite:///'):
                    sqlite_path = self.dsn[10:]  # 移除 'sqlite:///'
                    if sqlite_path.startswith('/'):
                        sqlite_db_path = sqlite_path
                    else:
                        # 相对路径，相对于项目根目录
                        sqlite_db_path = os.path.abspath(os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), sqlite_path))
                else:
                    sqlite_db_path = os.path.abspath(os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'data', 'aichat.db'))
            
                # 确保数据库目录存在
                db_dir = os.path.dirname(sqlite_db_path)
                if not os.path.exists(db_dir):
                    os.makedirs(db_dir, exist_ok=True)
                
                logger.info(f"尝试连接到SQLite数据库: {sqlite_db_path}")
                
                # 使用:memory:作为备选，避免路径编码问题
                try:
                    self.connection = sqlite3.connect(
                        sqlite_db_path,
                        detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
                    )
                except Exception:
                    logger.warning("SQLite文件连接失败，使用内存数据库")
                    self.connection = sqlite3.connect(
                        ':memory:',
                        detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
                    )
                
                # 启用外键约束
                self.connection.execute("PRAGMA foreign_keys = ON")
                
                # 启用WAL模式以提高并发性能
                self.connection.execute("PRAGMA journal_mode=WAL")
                
                # 使用row factory来模拟DictCursor
                def dict_factory(cursor, row):
                    d = {}
                    for idx, col in enumerate(cursor.description):
                        d[col[0]] = row[idx]
                    return d
                self.connection.row_factory = dict_factory
                logger.info("成功连接到SQLite数据库，已启用WAL模式")
            else:
                # 尝试连接PostgreSQL
                try:
                    # 使用DSN直接连接
                    logger.info(f"尝试连接到PostgreSQL数据库")
                    self.connection = psycopg2.connect(self.dsn)
                    # 注册向量类型
                    register_vector(self.connection)
                    logger.info(f"成功连接到PostgreSQL数据库")
                except psycopg2.OperationalError as e:
                    # 如果DSN连接失败，尝试从DSN解析参数
                    logger.warning(f"PostgreSQL DSN连接失败: {str(e)}，尝试解析DSN参数")
                    try:
                        from urllib.parse import urlparse
                        parsed = urlparse(self.dsn)
                        user = parsed.username or 'postgres'
                        password = parsed.password or ''
                        host = parsed.hostname or 'localhost'
                        port = parsed.port or 5432
                        dbname = parsed.path[1:] if parsed.path else 'aichat'
                        
                        # 处理密码编码问题
                        if isinstance(password, bytes):
                            password = password.decode('utf-8', errors='replace')
                        else:
                            password = str(password)
                        
                        self.connection = psycopg2.connect(
                            host=host,
                            port=port,
                            dbname=dbname,
                            user=user,
                            password=password,
                            sslmode='disable'
                        )
                        # 注册向量类型
                        register_vector(self.connection)
                        logger.info(f"成功连接到PostgreSQL数据库")
                    except UnicodeDecodeError as decode_e:
                        logger.error(f"PostgreSQL连接编码错误: {str(decode_e)}")
                        logger.warning("切换到SQLite数据库")
                        # 回退到SQLite
                        self.db_type = 'sqlite'
                        self._connect()
                    except Exception as inner_e:
                        logger.error(f"PostgreSQL参数连接失败: {str(inner_e)}")
                        logger.warning("切换到SQLite数据库")
                        # 回退到SQLite
                        self.db_type = 'sqlite'
                        self._connect()
                except UnicodeDecodeError as e:
                    logger.error(f"PostgreSQL连接编码错误: {str(e)}")
                    logger.warning("切换到SQLite数据库")
                    # 回退到SQLite
                    self.db_type = 'sqlite'
                    self._connect()
                except Exception as e:
                    logger.error(f"PostgreSQL连接失败: {str(e)}")
                    logger.warning("切换到SQLite数据库")
                    # 回退到SQLite
                    self.db_type = 'sqlite'
                    self._connect()
        except Exception as e:
            logger.error(f"数据库连接失败: {str(e)}")
            logger.warning("使用内存模式代替数据库")
            # 创建一个模拟的连接对象，避免程序崩溃
            self._use_in_memory_mode()
            
    def _use_in_memory_mode(self):
        """使用内存模式代替数据库"""
        logger.info("启用内存模式")
        # 创建一个模拟的连接对象
        class MockConnection:
            def __init__(self):
                self.closed = False  # 添加closed属性，默认False表示连接未关闭
            
            def cursor(self, cursor_factory=None):
                class MockCursor:
                    def execute(self, *args, **kwargs):
                        pass
                    def fetchone(self):
                        # 在内存模式下，返回一个模拟的文档ID
                        return {'id': 1}
                    def fetchall(self):
                        return []
                    def close(self):
                        pass
                    def __enter__(self):
                        return self
                    def __exit__(self, exc_type, exc_val, exc_tb):
                        pass
                return MockCursor()
            
            def commit(self):
                pass
            
            def close(self):
                self.closed = True  # 关闭连接时设置closed为True
            
            def rollback(self):
                pass  # 添加rollback方法，用于异常处理
        
        self.connection = MockConnection()
        self.in_memory_mode = True
    
    def _initialize_tables(self):
        """初始化数据库表结构"""
        try:
            with self.get_cursor() as cursor:
                if self.db_type == 'sqlite':
                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS users (
                        id VARCHAR(255) PRIMARY KEY,
                        username TEXT UNIQUE NOT NULL,
                        email TEXT UNIQUE,
                        password_hash TEXT NOT NULL,
                        nickname TEXT,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)

                    # SQLite表结构
                    # 创建文档表
                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS documents (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        content TEXT NOT NULL,
                        metadata TEXT,
                        vector TEXT,  -- 存储为JSON字符串
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)
                    
                    
                    # 创建学习数据表
                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS learning_data (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        user_query TEXT NOT NULL,
                        assistant_response TEXT,
                        corrected_response TEXT,
                        feedback TEXT,
                        is_positive INTEGER,
                        metadata TEXT,
                        used_in_training INTEGER DEFAULT 0,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)
                    
                    # 创建性能指标表
                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS performance_metrics (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        session_id VARCHAR(255),
                        request_id VARCHAR(255),
                        model_name VARCHAR(255),
                        processing_time FLOAT,
                        prompt_tokens INTEGER,
                        completion_tokens INTEGER,
                        total_tokens INTEGER,
                        success INTEGER,
                        error_message TEXT,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)

                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS session_model_map (
                        session_id VARCHAR(255) PRIMARY KEY,
                        model VARCHAR(255) NOT NULL,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)
                    cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_model_updated_at ON session_model_map(updated_at)")

                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS excel_datasets (
                        id VARCHAR(255) PRIMARY KEY,
                        name TEXT NOT NULL,
                        filename TEXT NOT NULL,
                        filepath TEXT NOT NULL,
                        filesize INTEGER,
                        sheet_name TEXT,
                        rows INTEGER,
                        cols INTEGER,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)

                    # 确保messages表存在user_id列
                    cursor.execute("PRAGMA table_info(messages)")
                    columns = {row['name'] if isinstance(row, dict) else row[1] for row in cursor.fetchall()}
                    if 'user_id' not in columns:
                        cursor.execute("ALTER TABLE messages ADD COLUMN user_id VARCHAR(255) DEFAULT 'default'")
                    cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_user_id ON messages(user_id)")
                else:
                    # PostgreSQL表结构
                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS users (
                        id VARCHAR(255) PRIMARY KEY,
                        username VARCHAR(255) UNIQUE NOT NULL,
                        email VARCHAR(255) UNIQUE,
                        password_hash TEXT NOT NULL,
                        nickname VARCHAR(255),
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)

                    cursor.execute("""
                    CREATE TABLE IF NOT EXISTS excel_datasets (
                        id VARCHAR(255) PRIMARY KEY,
                        name VARCHAR(255) NOT NULL,
                        filename VARCHAR(255) NOT NULL,
                        filepath TEXT NOT NULL,
                        filesize BIGINT,
                        sheet_name VARCHAR(255),
                        rows INTEGER,
                        cols INTEGER,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                    )
                    """)
                self._ensure_user_identifiers(cursor)
        except Exception as e:
            logger.error(f"初始化数据库表失败: {e}")
            raise

    def _ensure_user_identifiers(self, cursor):
        """确保users表具备username列及相关索引。"""
        if cursor is None:
            return

        try:
            if self.db_type == 'sqlite':
                cursor.execute("PRAGMA table_info(users)")
                rows = cursor.fetchall() or []
                columns = {row["name"] if isinstance(row, dict) else row[1] for row in rows}
            else:
                cursor.execute(
                    """
                    SELECT column_name
                    FROM information_schema.columns
                    WHERE table_name = 'users'
                      AND table_schema = 'public'
                    """
                )
                rows = cursor.fetchall() or []
                columns = {row["column_name"] if isinstance(row, dict) else row[0] for row in rows}

            if "username" not in columns:
                if self.db_type == "sqlite":
                    cursor.execute("ALTER TABLE users ADD COLUMN username TEXT")
                else:
                    cursor.execute("ALTER TABLE users ADD COLUMN username VARCHAR(255)")

            cursor.execute(
                """
                UPDATE users
                SET username = COALESCE(NULLIF(username, ''), email, id)
                WHERE username IS NULL OR username = ''
                """
            )
        except Exception as exc:
            logger.warning(f"保证users.username列存在失败: {exc}")

        # 确保索引存在，即便上面失败也不要阻塞启动
        try:
            cursor.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_users_username ON users(username)")
        except Exception as exc:
            logger.debug(f"创建idx_users_username失败: {exc}")
        try:
            cursor.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_users_email ON users(email)")
        except Exception as exc:
            logger.debug(f"创建idx_users_email失败: {exc}")
            
        # PostgreSQL特有的表结构初始化
        if self.db_type != 'sqlite':
            try:
                # 检查是否安装了pgvector扩展
                cursor.execute("CREATE EXTENSION IF NOT EXISTS vector")
                
                # 创建文档表，用于存储文档和向量
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id SERIAL PRIMARY KEY,
                    content TEXT NOT NULL,
                    metadata JSONB,
                    vector VECTOR(%s),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """, (self.vector_dimension,))
                
                # 创建向量索引
                cursor.execute("CREATE INDEX IF NOT EXISTS vector_idx ON documents USING ivfflat (vector)")
                
                
                # 创建学习数据表
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS learning_data (
                    id SERIAL PRIMARY KEY,
                    user_query TEXT NOT NULL,
                    assistant_response TEXT,
                    corrected_response TEXT,
                    feedback TEXT,
                    is_positive BOOLEAN,
                    metadata JSONB,
                    used_in_training BOOLEAN DEFAULT FALSE,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
                
                # 创建性能指标表
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS performance_metrics (
                    id SERIAL PRIMARY KEY,
                    session_id VARCHAR(255),
                    request_id VARCHAR(255),
                    model_name VARCHAR(255),
                    processing_time FLOAT,
                    prompt_tokens INTEGER,
                    completion_tokens INTEGER,
                    total_tokens INTEGER,
                    success BOOLEAN,
                    error_message TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)

                cursor.execute("""
                CREATE TABLE IF NOT EXISTS session_model_map (
                    session_id VARCHAR(255) PRIMARY KEY,
                    model VARCHAR(255) NOT NULL,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_model_updated_at ON session_model_map(updated_at)")
                cursor.execute("ALTER TABLE messages ADD COLUMN IF NOT EXISTS user_id VARCHAR(255) DEFAULT 'default'")
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_user_id ON messages(user_id)")
            except Exception as e:
                logger.warning(f"PostgreSQL表结构初始化部分失败: {str(e)}")
                
        try:
            self.connection.commit()
            logger.info("数据库表结构初始化成功")
        except Exception as e:
            if hasattr(self.connection, 'rollback'):
                self.connection.rollback()
            logger.error(f"数据库表结构初始化失败: {str(e)}")
            raise
    
    def commit(self):
        """提交当前事务"""
        if not self.in_memory_mode and self.connection and hasattr(self.connection, 'commit'):
            try:
                self.connection.commit()
                logger.debug("数据库事务提交成功")
            except Exception as e:
                logger.error(f"数据库事务提交失败: {str(e)}")
                raise
                
    def rollback(self):
        """回滚当前事务"""
        if not self.in_memory_mode and self.connection and hasattr(self.connection, 'rollback'):
            try:
                self.connection.rollback()
                logger.debug("数据库事务回滚成功")
            except Exception as e:
                logger.error(f"数据库事务回滚失败: {str(e)}")
                # 回滚失败不应阻止后续操作，只记录错误
                pass
                
    @contextmanager
    def get_cursor(self, autocommit=False, dict_cursor=True):
        """获取数据库游标上下文管理器，确保正确的事务管理和连接处理
        
        Args:
            autocommit: 是否自动提交事务，默认为False
            dict_cursor: 是否返回字典格式的游标，默认为True
            
        Yields:
            Cursor对象
        """
        # 如果是内存模式，直接返回模拟游标
        if self.in_memory_mode:
            yield self._mock_cursor
            return
            
        # 检查连接是否有效并重新连接
        if not self.connection:
            self._connect()
        else:
            # 检查SQLite连接是否有效
            if self.db_type == 'sqlite':
                try:
                    self.connection.execute("SELECT 1")
                except Exception:
                    self._connect()
            # 检查PostgreSQL连接是否有效
            elif hasattr(self.connection, 'closed') and self.connection.closed:
                self._connect()
        
        cursor = None
        original_row_factory = None  # 初始化变量，确保在finally块中可用
        
        try:
            # 设置autocommit
            if hasattr(self.connection, 'autocommit'):
                self.connection.autocommit = autocommit
            
            # 获取游标
            if self.db_type == 'sqlite':
                # 保存原始row_factory
                original_row_factory = self.connection.row_factory
                
                # 根据dict_cursor参数设置row_factory
                if dict_cursor:
                    # 设置字典工厂
                    def dict_factory(cursor, row):
                        d = {}
                        for idx, col in enumerate(cursor.description):
                            d[col[0]] = row[idx]
                        return d
                    self.connection.row_factory = dict_factory
                else:
                    # 使用默认row_factory（None）
                    self.connection.row_factory = None
                
                cursor = self.connection.cursor()
            else:
                if dict_cursor:
                    cursor = self.connection.cursor(cursor_factory=DictCursor)
                else:
                    cursor = self.connection.cursor()
            
            yield cursor
            
            # 如果不是自动提交，确保事务提交
            if not autocommit:
                self.commit()
        except Exception as e:
            logger.error(f"数据库操作失败: {str(e)}")
            # 发生错误时尝试回滚
            if not autocommit:
                try:
                    self.rollback()
                except Exception:
                    pass
            raise
        finally:
            # 确保游标被关闭
            if cursor and hasattr(cursor, 'close'):
                try:
                    cursor.close()
                except Exception:
                    pass
            
            # 对于SQLite，恢复原始row_factory
            if self.db_type == 'sqlite' and original_row_factory is not None:
                self.connection.row_factory = original_row_factory
    
    def close(self):
        """关闭数据库连接，确保在关闭前完成所有事务处理"""
        if self.connection:
            try:
                # 尝试提交可能未完成的事务
                try:
                    self.commit()
                except Exception:
                    # 如果提交失败，尝试回滚
                    try:
                        self.rollback()
                    except Exception:
                        pass
                # 关闭连接
                self.connection.close()
                logger.info("数据库连接已关闭")
            except Exception as e:
                logger.warning(f"关闭数据库连接时发生错误: {str(e)}")
    
    def insert_document(self, content, metadata, vector):
        """插入文档和向量"""
        try:
            with self.get_cursor() as cursor:
                if self.db_type == 'sqlite':
                    # SQLite需要将字典和向量转换为JSON字符串
                    metadata_str = json.dumps(metadata) if metadata else None
                    vector_str = json.dumps(vector.tolist()) if hasattr(vector, 'tolist') else json.dumps(vector)
                    cursor.execute(
                        "INSERT INTO documents (content, metadata, vector) VALUES (?, ?, ?)",
                        (content, metadata_str, vector_str)
                    )
                    document_id = cursor.lastrowid
                else:
                    # PostgreSQL直接支持JSONB和向量
                    cursor.execute(
                        "INSERT INTO documents (content, metadata, vector) VALUES (%s, %s, %s) RETURNING id",
                        (content, metadata, vector)
                    )
                    document_id = cursor.fetchone()['id']
                self.connection.commit()
                return document_id
        except Exception as e:
            if hasattr(self.connection, 'rollback'):
                self.connection.rollback()
            logger.error(f"插入文档失败: {str(e)}")
            raise
    
    def search_documents(self, query_vector, top_k=5, score_threshold=0.5):
        """基于向量相似度搜索文档"""
        try:
            with self.get_cursor() as cursor:
                if self.db_type == 'sqlite':
                    # SQLite简化版本，返回所有文档（实际项目中可能需要实现简单的相似度计算）
                    cursor.execute(
                        "SELECT id, content, metadata, 1.0 AS similarity FROM documents LIMIT ?",
                        (top_k,)
                    )
                else:
                    # PostgreSQL版本，使用向量相似度
                    cursor.execute(
                        """
                        SELECT id, content, metadata, 1 - (vector <=> %s) AS similarity 
                        FROM documents 
                        WHERE 1 - (vector <=> %s) > %s
                        ORDER BY vector <=> %s 
                        LIMIT %s
                        """,
                        (query_vector, query_vector, score_threshold, query_vector, top_k)
                    )
                results = cursor.fetchall()
                return results
        except Exception as e:
            logger.error(f"搜索文档失败: {str(e)}")
            raise
    
    def _update_session_timestamp(self, session_id):
        """更新会话的最后更新时间
        
        Args:
            session_id: 会话ID
        """
        from datetime import datetime
        # 对于SQLite，使用非字典格式的游标以避免类型不匹配
        use_dict_cursor = self.db_type != 'sqlite'
        
        with self.get_cursor(dict_cursor=use_dict_cursor) as cursor:
            current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            if self.db_type == 'sqlite':
                cursor.execute(
                    """UPDATE sessions SET updated_at = ? WHERE id = ?""",
                    (current_time, session_id)
                )
            else:
                # PostgreSQL版本
                cursor.execute(
                    """UPDATE sessions SET updated_at = %s WHERE id = %s""",
                    (current_time, session_id)
                )

    def _get_session_user_id(self, session_id):
        """获取会话所属用户ID."""
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute("SELECT user_id FROM sessions WHERE id = ?", (session_id,))
            else:
                cursor.execute("SELECT user_id FROM sessions WHERE id = %s", (session_id,))
            row = cursor.fetchone()
            if not row:
                return None
            if isinstance(row, dict):
                return row.get("user_id")
            return row[0]

    def get_user_by_username(self, username):
        """根据邮箱获取用户."""
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute(
                    "SELECT id, username, password_hash, nickname, created_at FROM users WHERE username = ?",
                    (username,),
                )
            else:
                cursor.execute(
                    "SELECT id, username, password_hash, nickname, created_at FROM users WHERE username = %s",
                    (username,),
                )
            return cursor.fetchone()

    def get_user_by_id(self, user_id):
        """根据ID获取用户."""
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute(
                    "SELECT id, username, password_hash, nickname, created_at FROM users WHERE id = ?",
                    (user_id,),
                )
            else:
                cursor.execute(
                    "SELECT id, username, password_hash, nickname, created_at FROM users WHERE id = %s",
                    (user_id,),
                )
        return cursor.fetchone()

    def get_session_by_id(self, session_id):
        """根据ID获取会话."""
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute(
                    "SELECT id, user_id, title, created_at, updated_at FROM sessions WHERE id = ?",
                    (session_id,),
                )
            else:
                cursor.execute(
                    "SELECT id, user_id, title, created_at, updated_at FROM sessions WHERE id = %s",
                    (session_id,),
                )
            row = cursor.fetchone()
            if not row:
                return None
            if isinstance(row, dict):
                return row
            return {
                "id": row[0],
                "user_id": row[1],
                "title": row[2] if len(row) > 2 else None,
                "created_at": row[3] if len(row) > 3 else None,
                "updated_at": row[4] if len(row) > 4 else None,
            }

    def assign_messages_user(self, session_id, user_id):
        """为会话下的消息设置user_id（仅针对缺失记录）。"""
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute(
                    """
                    UPDATE messages
                    SET user_id = ?
                    WHERE session_id = ?
                      AND (user_id IS NULL OR user_id = '' OR user_id = 'default')
                    """,
                    (user_id, session_id),
                )
            else:
                cursor.execute(
                    """
                    UPDATE messages
                    SET user_id = %s
                    WHERE session_id = %s
                      AND (user_id IS NULL OR user_id = '' OR user_id = 'default')
                    """,
                    (user_id, session_id),
                )

    def claim_session_owner(self, session_id, user_id):
        """若会话尚未绑定用户，则占用归属。"""
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute(
                    """
                    UPDATE sessions
                    SET user_id = ?
                    WHERE id = ?
                      AND (user_id IS NULL OR user_id = '' OR user_id = 'default')
                    """,
                    (user_id, session_id),
                )
            else:
                cursor.execute(
                    """
                    UPDATE sessions
                    SET user_id = %s
                    WHERE id = %s
                      AND (user_id IS NULL OR user_id = '' OR user_id = 'default')
                    """,
                    (user_id, session_id),
                )
            updated = cursor.rowcount or 0
        if updated:
            self.assign_messages_user(session_id, user_id)
        return bool(updated)

    def create_user(self, username, password_hash, nickname=None):
        """创建新用户."""
        user_id = str(uuid.uuid4())
        from datetime import datetime

        created_at = datetime.utcnow()
        email_value = username  # 兼容旧字段，默认与用户名一致
        with self.get_cursor() as cursor:
            if self.db_type == 'sqlite':
                cursor.execute(
                    """
                    INSERT INTO users (id, username, email, password_hash, nickname, created_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                    """,
                    (user_id, username, email_value, password_hash, nickname, created_at),
                )
            else:
                cursor.execute(
                    """
                    INSERT INTO users (id, username, email, password_hash, nickname, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    """,
                    (user_id, username, email_value, password_hash, nickname, created_at),
                )
        return {
            "id": user_id,
            "username": username,
            "password_hash": password_hash,
            "nickname": nickname,
            "created_at": created_at,
        }
    
    def insert_session_message(self, session_id, role, content, created_at=None, metadata=None, rag_context=None, message_id=None, user_id=None):
        """插入会话消息到messages表
        
        Args:
            session_id: 会话ID
            role: 角色（user, assistant, system等）
            content: 消息内容
            created_at: 创建时间，默认为None（自动生成）
            metadata: 元数据，默认为None
            rag_context: RAG上下文，默认为None
            message_id: 消息ID，默认为None（自动生成）
            
        Returns:
            插入的记录ID
        """
        # 对于SQLite，使用非字典格式的游标以避免lastrowid问题
        use_dict_cursor = self.db_type != 'sqlite'
        
        with self.get_cursor(dict_cursor=use_dict_cursor) as cursor:
            if message_id is None:
                import uuid
                message_id = str(uuid.uuid4())
            
            if created_at is None:
                from datetime import datetime
                created_at = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            
            metadata_payload = metadata.copy() if isinstance(metadata, dict) else {}
            if message_id:
                metadata_payload.setdefault("message_id", str(message_id))
            if rag_context is not None:
                metadata_payload["rag_context"] = rag_context
            metadata_text = json.dumps(metadata_payload) if metadata_payload else None
            if not user_id:
                user_id = self._get_session_user_id(session_id) or "default"

            if self.db_type == 'sqlite':
                cursor.execute(
                    f"""
                    INSERT INTO messages (session_id, user_id, role, content, meta, created_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                    """,
                    (
                        session_id,
                        user_id,
                        role,
                        content,
                        metadata_text,
                        created_at,
                    )
                )
                row_id = cursor.lastrowid
            else:
                cursor.execute(
                    """
                    INSERT INTO messages (session_id, user_id, role, content, meta, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s) RETURNING id
                    """,
                    (
                        session_id,
                        user_id,
                        role,
                        content,
                        metadata_payload or None,
                        created_at
                    )
                )
                result = cursor.fetchone()
                row_id = result['id'] if result else None
            
            try:
                self._update_session_timestamp(session_id)
            except Exception as e:
                logger.warning(f"更新会话时间戳失败: {str(e)}")
            
            return row_id
    
    def get_session_messages(self, session_id, limit=20, offset=0):
        """获取会话消息记录
        
        Args:
            session_id: 会话ID
            limit: 返回记录数限制
            offset: 偏移量
            
        Returns:
            聊天记录列表
        """
        try:
            # 对于SQLite，使用非字典格式的游标以避免row_factory问题
            use_dict_cursor = self.db_type != 'sqlite'
            
            with self.get_cursor(dict_cursor=use_dict_cursor) as cursor:
                if self.db_type == 'sqlite':
                    cursor.execute(
                        """
                        SELECT id, session_id, role, content, created_at, meta
                        FROM messages
                        WHERE session_id = ?
                        ORDER BY created_at ASC, id ASC
                        LIMIT ? OFFSET ?
                        """,
                        (session_id, limit, offset)
                    )
                else:
                    cursor.execute(
                        """
                        SELECT id, session_id, role, content, created_at, meta
                        FROM messages
                        WHERE session_id = %s
                        ORDER BY created_at ASC, id ASC
                        LIMIT %s OFFSET %s
                        """,
                        (session_id, limit, offset)
                    )
                
                results = cursor.fetchall()
                formatted_results = []
                for row in results:
                    if isinstance(row, dict):
                        metadata_value = row.get('meta')
                        if isinstance(metadata_value, str):
                            metadata_value = json.loads(metadata_value) if metadata_value else {}
                        formatted_row = {
                            'id': str(row.get('id')) if row.get('id') is not None else None,
                            'session_id': str(row.get('session_id')) if row.get('session_id') is not None else None,
                            'role': row.get('role'),
                            'content': row.get('content'),
                            'metadata': metadata_value or {},
                            'created_at': row.get('created_at')
                        }
                    else:
                        # SQLite返回的是元组，需要检查元组长度
                        if len(row) >= 6:
                            # 有meta字段的情况
                            metadata_value = row[5]
                            if isinstance(metadata_value, str):
                                metadata_value = json.loads(metadata_value) if metadata_value else {}
                            formatted_row = {
                                'id': str(row[0]) if row[0] is not None else None,
                                'session_id': str(row[1]) if row[1] is not None else None,
                                'role': row[2],
                                'content': row[3],
                                'metadata': metadata_value or {},
                                'created_at': row[4]
                            }
                        else:
                            # 没有meta字段的情况（只有5个元素）
                            formatted_row = {
                                'id': str(row[0]) if row[0] is not None else None,
                                'session_id': str(row[1]) if row[1] is not None else None,
                                'role': row[2],
                                'content': row[3],
                                'metadata': {},
                                'created_at': row[4]
                            }
                    if isinstance(formatted_row['metadata'], str):
                        formatted_row['metadata'] = json.loads(formatted_row['metadata'] or '{}')
                    rag_context = formatted_row['metadata'].get('rag_context') if formatted_row['metadata'] else None
                    formatted_row['rag_context'] = rag_context or []
                    message_identifier = formatted_row['metadata'].get('message_id') if formatted_row['metadata'] else None
                    if message_identifier:
                        formatted_row['message_id'] = message_identifier
                    formatted_results.append(formatted_row)
                
                return formatted_results
        except Exception as e:
            logger.error(f"获取聊天历史失败: {str(e)}")
            if hasattr(self, '_handle_error'):
                self._handle_error()
            return []
    
    def count_session_messages(self, session_id):
        """
        获取聊天历史记录总数
        
        Args:
            session_id: 会话ID
            
        Returns:
            消息总数
        """
        try:
            # 对于SQLite，使用非字典格式的游标以避免类型不匹配
            use_dict_cursor = self.db_type != 'sqlite'
            
            with self.get_cursor(dict_cursor=use_dict_cursor) as cursor:
                if self.db_type == 'sqlite':
                    cursor.execute("SELECT COUNT(*) AS total FROM messages WHERE session_id = ?", (session_id,))
                else:
                    cursor.execute("SELECT COUNT(*) AS total FROM messages WHERE session_id = %s", (session_id,))
                result = cursor.fetchone()
                if isinstance(result, dict):
                    return int(result.get('total') or result.get('count') or 0)
                # SQLite返回的是元组，只有一个值
                return int(result[0]) if result else 0
        except Exception as e:
            logger.error(f"获取聊天历史数量失败: {str(e)}")
            if hasattr(self, '_handle_error'):
                self._handle_error()
            return 0
    
    def delete_session_messages(self, session_id):
        """
        删除聊天历史记录
        
        Args:
            session_id: 会话ID
            
        Returns:
            删除的记录数
        """
        try:
            with self.get_cursor() as cursor:
                
                if self.db_type == 'sqlite':
                    cursor.execute(
                        """
                        DELETE FROM messages 
                        WHERE session_id = ?
                        """,
                        (session_id,)
                    )
                else:
                    cursor.execute(
                        """
                        DELETE FROM messages 
                        WHERE session_id = %s
                        """,
                        (session_id,)
                    )
                
                self.connection.commit()
                return cursor.rowcount
        except Exception as e:
            logger.error(f"删除聊天历史失败: {str(e)}")
            if hasattr(self.connection, 'rollback'):
                self.connection.rollback()
            if hasattr(self, '_handle_error'):
                self._handle_error()
            return 0

    def set_session_model(self, session_id: str, model: str):
        """为指定会话设置模型映射"""
        from datetime import datetime
        updated_at = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        try:
            with self.get_cursor() as cursor:
                if self.db_type == 'sqlite':
                    cursor.execute(
                        """
                        INSERT INTO session_model_map (session_id, model, updated_at)
                        VALUES (?, ?, ?)
                        ON CONFLICT(session_id) DO UPDATE SET
                            model = excluded.model,
                            updated_at = excluded.updated_at
                        """,
                        (session_id, model, updated_at)
                    )
                else:
                    cursor.execute(
                        """
                        INSERT INTO session_model_map (session_id, model, updated_at)
                        VALUES (%s, %s, %s)
                        ON CONFLICT (session_id) DO UPDATE SET
                            model = EXCLUDED.model,
                            updated_at = EXCLUDED.updated_at
                        """,
                        (session_id, model, updated_at)
                    )
        except Exception as e:
            logger.error(f"设置会话模型失败: {str(e)}")
            if hasattr(self, '_handle_error'):
                self._handle_error()
            raise

    def get_session_model(self, session_id: str) -> Optional[str]:
        """获取指定会话绑定的模型"""
        try:
            # 对于SQLite，使用非字典格式的游标以避免类型不匹配
            use_dict_cursor = self.db_type != 'sqlite'
            
            with self.get_cursor(dict_cursor=use_dict_cursor) as cursor:
                if self.db_type == 'sqlite':
                    cursor.execute(
                        "SELECT model FROM session_model_map WHERE session_id = ?",
                        (session_id,)
                    )
                else:
                    cursor.execute(
                        "SELECT model FROM session_model_map WHERE session_id = %s",
                        (session_id,)
                    )
                result = cursor.fetchone()
                if not result:
                    return None
                if isinstance(result, dict):
                    return result.get('model')
                # SQLite返回的是元组，只有一个值
                return result[0] if result else None
        except Exception as e:
            logger.error(f"获取会话模型失败: {str(e)}")
            if hasattr(self, '_handle_error'):
                self._handle_error()
            return None

    def delete_session_model(self, session_id: str):
        """删除指定会话的模型映射"""
        try:
            with self.get_cursor() as cursor:
                if self.db_type == 'sqlite':
                    cursor.execute(
                        "DELETE FROM session_model_map WHERE session_id = ?",
                        (session_id,)
                    )
                else:
                    cursor.execute(
                        "DELETE FROM session_model_map WHERE session_id = %s",
                        (session_id,)
                    )
                return cursor.rowcount
        except Exception as e:
            logger.error(f"删除会话模型映射失败: {str(e)}")
            if hasattr(self, '_handle_error'):
                self._handle_error()
            return 0
    
    def delete_vector(self, vector_id):
        """删除指定ID的向量"""
        if self.in_memory_mode:
            # 内存模式
            for i, vector in enumerate(self.memory_vectors):
                if vector['id'] == vector_id:
                    self.memory_vectors.pop(i)
                    return True
            return False
        
        try:
            with self.get_cursor() as cursor:
                cursor.execute(
                    "DELETE FROM documents WHERE id = %s",
                    (vector_id,)
                )
                self.connection.commit()
                return cursor.rowcount > 0
        except Exception as e:
            self.connection.rollback()
            logger.error(f"删除向量失败: {str(e)}")
            return False
    
    def count_vectors(self):
        """获取向量总数"""
        if self.in_memory_mode:
            return len(self.memory_vectors)
        
        try:
            with self.get_cursor() as cursor:
                cursor.execute("SELECT COUNT(*) FROM documents")
                count = cursor.fetchone()[0]
                return count
        except Exception as e:
            logger.error(f"获取向量总数失败: {str(e)}")
            return -1
    
    def create_tables(self):
        """创建数据库表结构（兼容vector_db_manager的方法名）"""
        self._initialize_tables()
    
    def insert_learning_data(self, user_query, assistant_response, corrected_response=None, feedback=None, is_positive=None, metadata=None):
        """插入学习数据"""
        try:
            with self.get_cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO learning_data (user_query, assistant_response, corrected_response, feedback, is_positive, metadata) 
                    VALUES (%s, %s, %s, %s, %s, %s)
                    """,
                    (user_query, assistant_response, corrected_response, feedback, is_positive, metadata)
                )
                self.connection.commit()
        except Exception as e:
            self.connection.rollback()
            logger.error(f"插入学习数据失败: {str(e)}")
            raise
    
    def get_unused_learning_data(self, limit=None):
        """获取未用于训练的学习数据"""
        try:
            with self.get_cursor() as cursor:
                query = "SELECT id, user_query, assistant_response, corrected_response FROM learning_data WHERE used_in_training = false"
                params = ()
                
                if limit:
                    query += " LIMIT %s"
                    params = (limit,)
                
                cursor.execute(query, params)
                results = cursor.fetchall()
                return results
        except Exception as e:
            logger.error(f"获取学习数据失败: {str(e)}")
            raise
    
    def mark_learning_data_as_used(self, ids):
        """标记学习数据为已使用"""
        try:
            with self.get_cursor() as cursor:
                cursor.execute(
                    "UPDATE learning_data SET used_in_training = true WHERE id = ANY(%s)",
                    (ids,)
                )
                self.connection.commit()
        except Exception as e:
            self.connection.rollback()
            logger.error(f"标记学习数据失败: {str(e)}")
            raise
    
    def insert_performance_metrics(self, metrics):
        """插入性能指标"""
        try:
            with self.get_cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO performance_metrics (
                        session_id, request_id, model_name, processing_time, prompt_tokens, 
                        completion_tokens, total_tokens, success, error_message
                    ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        metrics.get('session_id'),
                        metrics.get('request_id'),
                        metrics.get('model_name'),
                        metrics.get('processing_time'),
                        metrics.get('prompt_tokens'),
                        metrics.get('completion_tokens'),
                        metrics.get('total_tokens'),
                        metrics.get('success'),
                        metrics.get('error_message')
                    )
                )
                self.connection.commit()
        except Exception as e:
            self.connection.rollback()
            logger.error(f"插入性能指标失败: {str(e)}")
            raise

# 全局数据库管理器实例
db_manager = DatabaseManager()
